import json
from typing import List
from os.path import abspath, join
from backend.experiment.clue.c3.scheme1.implement import C3Implement
from backend.experiment.tokenizer import tokenizer
from backend.experiment.model import vocab_name
from backend.experiment.framework.shots_prompt.implement import ShotMode
from backend.experiment.framework.task import PerplexityChoiceMakingTask

test_dev = True
insert_start_token = True
only_calculate_last = True
seed = 42
task_handler = C3Implement.get_task_handler()


class C3Sample:
    def __init__(
            self, question: str, candidates: List[str], answer: str,
            text: str, task
    ):
        self.__question = question
        self.__candidates = candidates
        self.__answer = answer
        self.__text = text
        self.__task = task

    @property
    def question(self) -> str:
        return self.__question

    @property
    def candidates(self) -> List[str]:
        return self.__candidates

    @property
    def answer(self) -> str:
        return self.__answer

    @property
    def text(self) -> str:
        return self.__text

    @property
    def task(self) -> PerplexityChoiceMakingTask:
        return self.__task


def initialize() -> List[C3Sample]:
    with open(abspath(
            join(__file__, '../../experiment/clue/c3/d-dev.json')),
            'r', encoding='utf8') as file:
        records = json.loads(file.read())
    with open(abspath(
            join(__file__, '../../experiment/clue/c3/m-dev.json')),
            'r', encoding='utf8') as file:
        records.extend(json.loads(file.read()))

    cache_filename = f'{vocab_name}_one_shot_{seed}'
    if insert_start_token:
        cache_filename += '_start_token'

    if only_calculate_last:
        cache_filename += '_only_last_shot'

    test_tasks = C3Implement(
        tokenizer=tokenizer,
        test_dev=test_dev,
        seed=seed
    ).get_tasks(
        mode=ShotMode.FEW_SHOT,
        cache_filename=cache_filename + '.json',
        insert_start_token=insert_start_token,
        only_calculate_last=only_calculate_last
    )

    _samples = []
    i = 0
    for each in records:
        text = '\n'.join(each[0])
        for q in each[1]:
            if test_tasks[i] is not None:
                _samples.append(
                    C3Sample(
                        question=q['question'],
                        candidates=q['choice'],
                        answer=q['answer'],
                        text=text,
                        task=test_tasks[i]
                    ))
            i += 1
    return _samples


samples = initialize()
